package org.chmodke.proxy.socks.server;

import org.chmodke.proxy.socks.util.C;
import org.chmodke.proxy.socks.util.P;
import org.chmodke.proxy.socks.util.U;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author kehao
 * @version 1.0
 * @date 2019/3/15 10:19
 */
public class SocksServerThread implements Runnable {
    private static final Logger logger = LoggerFactory.getLogger(SocksServerThread.class);
    /**
     * 来源的代理socket
     */
    private final Socket socket;
    /**
     * 是否开启socks4代理
     */
    private final boolean openSock4;
    /**
     * 是否开启socks5代理
     */
    private final boolean openSock5;
    /**
     * socks5代理的登录用户名，如果 不为空表示需要登录验证
     */
    private final String user;
    /**
     * socks5代理的登录密码，
     */
    private final String pwd;
    /**
     * socks是否需要进行登录验证
     */
    private final boolean socksNeedLogin;

    /**
     * @param socket    来源的代理socket
     * @param openSock4 是否开启socks4代理
     * @param openSock5 是否开启socks5代理
     * @param user      socks5代理的登录用户名，如果 不为空表示需要登录验证
     * @param pwd       socks5代理的登录密码，
     */
    protected SocksServerThread(Socket socket, boolean openSock4, boolean openSock5, String user, String pwd) {
        this.socket = socket;
        this.openSock4 = openSock4;
        this.openSock5 = openSock5;
        this.user = user;
        this.pwd = pwd;
        this.socksNeedLogin = null != user;
    }

    /**
     * 数据交换.主要用于tcp协议的交换
     *
     * @param in    输入流
     * @param out   输出流
     * @param alive 链路是否存活
     */
    protected static void transfer(final InputStream in, final OutputStream out, final AtomicBoolean alive) {
        Transfer transfer = new Transfer(in, out, alive);
        C.PROXY.execute(transfer);
    }

    public static void startServer(int port, boolean openSock4, boolean openSock5, String user, String pwd)
            throws IOException {
        logger.info("{} port:{}, openSock4:{}, openSock5:{}, user:{}, pwd:xxxxxx.", C.MAIN_LOG, port, openSock4,
                openSock5, user);
        ServerSocket ss = new ServerSocket(port);
        logger.info("{} Socks server port: {} listening...", C.MAIN_LOG, port);
        Socket socket;
        while (null != (socket = ss.accept())) {
            SocksServerThread serverThread = new SocksServerThread(socket, openSock4, openSock5, user, pwd);
            C.ACCEPT.execute(serverThread);
        }
        ss.close();
    }

    @Override
    public void run() {
        // 获取来源的地址用于日志打印使用
        String addr = socket.getRemoteSocketAddress().toString();
        logger.info("{} start process one socket: {}", C.PROXY_LOG, addr);
        // 声明流
        InputStream srcSocksIn, destSocksIn;
        OutputStream srcSocksOut, destSocksOut;
        Socket destSocket = null;
        try {
            srcSocksIn = socket.getInputStream();
            srcSocksOut = socket.getOutputStream();

            // 获取协议头。取代理的类型，只有 4，5。
            byte version = (byte) srcSocksIn.read();
            if (version > -1) {
                if ((openSock4 && P.isV4(version))) {
                    // 如果开启代理4，并以socks4协议请求
                    destSocket = handleSock4(srcSocksIn, srcSocksOut);
                } else if ((openSock5 && P.isV5(version))) {
                    // 如果开启代理5，并以socks5协议请求
                    destSocket = handleSock5(srcSocksIn, srcSocksOut);
                } else {
                    // 非socks 4 ,5 协议的请求
                    logger.error("{} not socks{} proxy, openSock4:{} openSock5:{}.", P.version(version), C.MAIN_LOG,
                            openSock4,
                            openSock5);
                }
                if (null != destSocket) {
                    destSocksIn = destSocket.getInputStream();
                    destSocksOut = destSocket.getOutputStream();
                    // 交换流数据
                    AtomicBoolean alive = new AtomicBoolean(true);
                    transfer(srcSocksIn, destSocksOut, alive);
                    transfer(destSocksIn, srcSocksOut, alive);
                }
            } else {
                logger.error("{} socks protocol is error : {}", C.PROXY_LOG, version);
            }
        } catch (Exception e) {
            logger.error("{} exception:{}", C.PROXY_LOG, e);
        }
    }


    /**
     * sock4代理的头处理
     *
     * @param srcInput  客户端输入流
     * @param srcOutput 客户端输出流
     * @return 代理目标Socket
     * @throws IOException IOException
     */
    private Socket handleSock4(InputStream srcInput, OutputStream srcOutput) throws IOException {
        byte cmd = (byte) srcInput.read();
        byte[] ports = new byte[2];
        srcInput.read(ports);
        short destPort = ByteBuffer.wrap(ports).getShort();
        byte[] destHost = P.V5.getHost(P.V5.ADDRESS_TYPE_IPV4, srcInput);
        String destIp = InetAddress.getByAddress(destHost).getHostAddress();
        // ---------------------
        ByteBuffer respBuffer = ByteBuffer.allocate(8);
        respBuffer.put(P.VER4);

        Object middleSocket = null;
        try {
            switch (cmd) {
                case P.V4.CD_CONNECT:
                    middleSocket = new Socket(InetAddress.getByAddress(destHost), destPort);
                    // 代理成功
                    respBuffer.put(P.V4.RESP_ACCEPT_FORWARD);
                    break;
                case P.V4.CD_BIND:
                    middleSocket = new ServerSocket(destPort);
                    // 代理成功
                    respBuffer.put(P.V4.RESP_ACCEPT_FORWARD);
                    break;
                default:
                    respBuffer.put(P.V4.RESP_REJECT_FORWARD);
                    break;
            }
            logger.info("{} connected to {}:{}", C.PROXY_LOG, destIp, destPort);
        } catch (Exception e) {
            // 代理失败.
            logger.error("{} connect exception  {}:{}", C.PROXY_LOG, destIp, destPort);
            respBuffer.put(P.V4.RESP_REJECT_FORWARD);
        }
        respBuffer.put(socket.getLocalAddress().getAddress());
        short localPort = (short) (socket.getLocalPort() & P.FFFF);
        respBuffer.putShort(localPort);
        srcOutput.write(respBuffer.array());
        srcOutput.flush();
        if (null != middleSocket && P.V4.CD_BIND == cmd) {
            ServerSocket ss = (ServerSocket) middleSocket;
            try {
                middleSocket = ss.accept();
            } catch (IOException e) {
            } finally {
                U.closeIo(ss);
            }
        }
        return (Socket) middleSocket;
    }


    /**
     * socks5认证
     *
     * @param user 用户名
     * @param pwd  密码
     * @return 认证结果
     */
    private boolean login(String user, String pwd) {
        return null != user && user.trim().equals(this.user) && null != pwd && pwd.trim().equals(this.pwd);
    }

    /**
     * sock5代理头处理
     *
     * @param srcInput  客户端输入流
     * @param srcOutput 客户端输出流
     * @return 代理目标Socket
     * @throws IOException IOException
     */
    private Socket handleSock5(InputStream srcInput, OutputStream srcOutput) throws IOException {
        // 读取认证方式数量
        byte count = (byte) srcInput.read();

        //读取认证方式列表
        byte[] methods = new byte[count];
        srcInput.read(methods);
        // ----------------
        byte[] resp = new byte[2];
        resp[0] = P.VER5;
        if (socksNeedLogin && P.V5.isPwdAuth(methods)) {
            //响应密码认证
            resp[1] = P.V5.AUTH_PWD;
        } else if (!socksNeedLogin && P.V5.noneAuth(methods)) {
            //响应不用认证
            resp[1] = P.V5.AUTH_NONE;
        } else {
            // 响应无匹配的认证方式
            resp[1] = (byte) (P.V5.AUTH_NOT_MATCH & P.FF);
        }
        srcOutput.write(resp);
        srcOutput.flush();
        if (resp[1] == P.V5.AUTH_NOT_MATCH) {
            return null;
        }
        // ----------------

        boolean isLogin = false;
        if (socksNeedLogin && P.V5.isPwdAuth(methods)) {
            // 处理登录.
            int ver = srcInput.read();
            if (P.V5.SUB_AUTH_VER == ver) {
                int uLen = srcInput.read();
                byte[] nameArray = new byte[uLen];
                srcInput.read(nameArray);
                String user = new String(nameArray);

                int pLen = srcInput.read();
                byte[] pwdArray = new byte[pLen];
                srcInput.read(pwdArray);
                String pwd = new String(pwdArray);

                if (login(user, pwd)) {
                    // 登录成功
                    isLogin = true;
                    resp = new byte[]{0x01, P.V5.AUTH_SUCCEED};
                    logger.info("{} {} login success !", C.PROXY_LOG, user);
                } else {
                    // 登录失败
                    resp = new byte[]{0x01, P.V5.AUTH_FAIL};
                    logger.warn("{} {}  login fail !", C.PROXY_LOG, user);
                }
                srcOutput.write(resp);
                srcOutput.flush();
                if (resp[1] != P.V5.AUTH_SUCCEED) {
                    return null;
                }
            }
        }
        // ----------------

        Object middleSocket = null;
        byte cmd = 0;
        if (!socksNeedLogin || isLogin) {
            // V|C|R|T
            byte[] head = new byte[4];
            srcInput.read(head);
            logger.debug("{} proxy header is>> {}", C.PROXY_LOG, Arrays.toString(head));
            cmd = head[1];
            byte[] destHost = P.V5.getHost(head[3], srcInput);
            String destIp = InetAddress.getByAddress(destHost).getHostAddress();
            byte[] ports = new byte[2];
            srcInput.read(ports);
            short destPort = ByteBuffer.wrap(ports).getShort();
            logger.info("{} connecting {}:{}", C.PROXY_LOG, destIp, destPort);

            ByteBuffer respBuffer = ByteBuffer.allocate(10);
            respBuffer.put(P.VER5);
            try {
                switch (cmd) {
                    case P.V5.CD_CONNECT:
                        middleSocket = new Socket(InetAddress.getByAddress(destHost), destPort);
                        respBuffer.put(P.V5.RESP_SUCCEED);
                        break;
                    case P.V5.CD_BIND:
                        middleSocket = new ServerSocket(destPort);
                        respBuffer.put(P.V5.RESP_SUCCEED);
                        break;
                    default:
                        middleSocket = null;
                        respBuffer.put(P.V5.RESP_UNSUPPORTED_COMMAND);
                        break;
                }
                logger.info("{} connected to {}:{}", C.PROXY_LOG, destIp, destPort);
            } catch (IOException e) {
                logger.error("{} connect exception  {}:{}", C.PROXY_LOG, destIp, destPort);
                respBuffer.put(P.V5.RESP_SERVER_FAULT);
                middleSocket = null;
            }
            respBuffer.put(P.INVALID);
            respBuffer.put(P.V5.ADDRESS_TYPE_IPV4);
            respBuffer.put(socket.getLocalAddress().getAddress());
            short localPort = (short) (socket.getLocalPort() & P.FFFF);
            respBuffer.putShort(localPort);
            resp = respBuffer.array();
        } else {
            // 登录失败
            resp = new byte[]{P.VER5, P.V5.AUTH_FAIL};
            logger.warn("{} socks server need login,but no login info.", C.PROXY_LOG);
        }
        srcOutput.write(resp);
        srcOutput.flush();
        if (null != middleSocket && P.V5.CD_BIND == cmd) {
            ServerSocket ss = (ServerSocket) middleSocket;
            try {
                middleSocket = ss.accept();
            } catch (IOException e) {
            } finally {
                U.closeIo(ss);
            }
        }
        return (Socket) middleSocket;
    }
}
